import re
from datetime import datetime
from dataclasses import dataclass

@dataclass
class DJIMetadata:
    frame_cnt: int
    timestamp: datetime
    focal_len: float
    latitude: float
    longitude: float
    altitude: float
    yaw: float
    pitch: float
    roll: float
    sensor_type: str  # 新增传感器类型标识

class DJISrtParser:
    def __init__(self):
        self.meta_pattern = re.compile(r"""
            FrameCnt:\s*(\d+).*?
            (\d{4}-\d{2}-\d{2}\s\d{2}:\d{2}:\d{2}\.\d{3}).*?
            focal_len:\s*([\d.]+).*?
            latitude:\s*(-?\d+\.\d+).*?
            longitude:\s*(-?\d+\.\d+).*?
            rel_alt:\s*([\d.]+).*?
            gb_yaw:\s*([-\d.]+).*?
            gb_pitch:\s*([-\d.]+).*?
            gb_roll:\s*([-\d.]+)
        """, re.DOTALL | re.VERBOSE)

    def parse(self, file_path):
        """解析SRT文件并返回元数据列表"""
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
            
        sensor_type = "Thermal" if "_T." in file_path else "Visible"
        metadata = []
        
        for match in re.finditer(self.meta_pattern, content):
            groups = match.groups()
            meta = DJIMetadata(
                frame_cnt=int(groups[0]),
                timestamp=datetime.strptime(groups[1], "%Y-%m-%d %H:%M:%S.%f"),
                focal_len=float(groups[2]),
                latitude=float(groups[3]),
                longitude=float(groups[4]),
                altitude=float(groups[5]),
                yaw=float(groups[6]),
                pitch=float(groups[7]),
                roll=float(groups[8]),
                sensor_type=sensor_type
            )
            metadata.append(meta)
        return metadata
    